/// Copyright (c) dexterdreeeam​
/// Original Post: https://zhuanlan.zhihu.com/p/460830735
#pragma once

#include <atomic>
#include <coroutine>
#include <iostream>
#include <memory>
#include <thread>

class task_fork {
 public:
  bool await_ready() { return false; }

  void await_suspend(std::coroutine_handle<> coro) {
    std::thread t([=] {
      std::cout << "New thread id " << std::this_thread::get_id() << std::endl;
      coro.resume();
    });
    t.detach();
  }

  void await_resume() {}
};

enum class task_context_state : int {
  Init,           // task init
  Complete,       // task complete by sub_thread
  Awaiting,       // there is coroutine awaiting me
  ParentCoroSet,  // parent coroutine handler is set, sub_thread could resume on
                  // parent remaining content
};

struct task_context {
  std::coroutine_handle<> _coro;
  std::coroutine_handle<> _parent_coro;
  std::atomic<task_context_state> _state;
};

class task {
 public:
  struct promise_type {
    promise_type() { std::cout << "Promise +" << std::endl; }

    ~promise_type() { std::cout << "Promise -" << std::endl; }

    auto initial_suspend() noexcept { return std::suspend_never(); }

    auto final_suspend() noexcept {
      task_context_state exp = task_context_state::Init;
      if (!_ctx->_state.compare_exchange_strong(exp,
                                                task_context_state::Complete)) {
        while (_ctx->_state.load() != task_context_state::ParentCoroSet) {
          std::this_thread::yield();
        }
        _ctx->_parent_coro.resume();
      }
      return std::suspend_never();
    }

    task get_return_object() {
      auto ctx = std::make_shared<task_context>();
      ctx->_coro = std::coroutine_handle<promise_type>::from_promise(*this);
      ctx->_parent_coro = nullptr;
      ctx->_state.store(task_context_state::Init);

      _ctx = ctx;
      return task(ctx);
    }

    void return_void() {}

    void unhandled_exception() noexcept { return; }

    std::shared_ptr<task_context> _ctx;
  };

  task() = delete;

  task(std::shared_ptr<task_context> ctx) : _ctx(ctx) {}

  task(const task& rhs) : _ctx(rhs._ctx) {}

  bool await_ready() {
    task_context_state exp = task_context_state::Init;
    if (_ctx->_state.compare_exchange_strong(exp,
                                             task_context_state::Awaiting)) {
      return false;
    } else {
      return true;
    }
  }

  void await_suspend(std::coroutine_handle<> coro) {
    _ctx->_parent_coro = coro;
    _ctx->_state.store(task_context_state::ParentCoroSet);
  }

  void await_resume() {}

 public:
  std::shared_ptr<task_context> _ctx;
};
